
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <unistd.h>
#include <fcntl.h>
#include <signal.h>
#include <inttypes.h>
#include <termios.h>
#include <poll.h>
#include <time.h>

#include "cutils.h"
#include "virtio.h"
#include "machine.h"
#include "riscv_cpu.h"

#define MAX_EXEC_CYCLE 500000
#define MB (1024 * 1024)
#define MAX_RAM_SIZE (128 * MB)
#define DEFAULT_DRIVE "./fs.img"

typedef struct {
    int stdin_fd;
    int stdout_fd;
    int console_esc_state;
    // BOOL resize_pending;
} STDIODevice;

static struct termios oldtty;

static int old_fd0_flags;

static STDIODevice global_stdio_device;

static void
term_exit () {
    tcsetattr (0, TCSANOW, &oldtty);
    fcntl (0, F_SETFL, old_fd0_flags);
}

static void
term_init (BOOL allow_ctrlc) {
    struct termios tty;
    memset (&tty, 0, sizeof (tty));
    memset (&global_stdio_device, 0 ,sizeof (global_stdio_device));
    global_stdio_device.stdin_fd = STDIN_FILENO;
    global_stdio_device.stdout_fd = STDOUT_FILENO;

    tcgetattr (0, &tty);
    oldtty = tty;
    old_fd0_flags = fcntl(0, F_GETFL);

    tty.c_iflag &= ~(IGNBRK|BRKINT|PARMRK|ISTRIP
                          |INLCR|IGNCR|ICRNL|IXON);
    tty.c_oflag |= OPOST;
    tty.c_lflag &= ~(ECHO|ECHONL|ICANON|IEXTEN);
    if (!allow_ctrlc)
        tty.c_lflag &= ~ISIG;
    tty.c_cflag &= ~(CSIZE|PARENB);
    tty.c_cflag |= CS8;
    tty.c_cc[VMIN] = 1;
    tty.c_cc[VTIME] = 0;

    tcsetattr (0, TCSANOW, &tty);

    atexit(term_exit);
}

static void
console_write (void *opaque, const uint8_t *buf, int len) {
    STDIODevice *s = opaque;
    write (s->stdout_fd, buf, len);
}

static int
console_read (void * opaque, uint8_t *buf, int len) {
    STDIODevice *s = opaque;
    int ret, i, j;
    uint8_t ch;
    if (len <= 0)
        return 0;
    
    ret = read (s->stdin_fd, buf, len);
    if (ret < 0)
        return 0;
    if (ret == 0) {
        /* EOF */
        exit (EXIT_FAILURE);
    }

    j = 0;
    for (i = 0; i < ret; i++) {
        ch = buf[i];
        if (s->console_esc_state) {
            s->console_esc_state = 0;
            switch (ch)
            {
            case 'x':
                printf ("Terminated.\n");
                exit (EXIT_SUCCESS);
                break;
            case 'h':
                printf("\n"
                       "C-a h   print this help\n"
                       "C-a x   exit vm\n"
                       "C-a C-a send C-a\n"
                        );
                break;
            case 1:
                goto output_char;
            default:
                break;
            }
        } else {
            if (ch == 1) {
                s->console_esc_state = 1;
            } else {
                output_char:
                buf[j++] = ch;
            }
        }
    }
    return j;
}

CharacterDevice *
console_init () {
    CharacterDevice *s;
    s = mallocz (sizeof (*s));
    s->opaque = &global_stdio_device;
    s->read_data = console_read;
    s->write_data = console_write;
    return s;
}

typedef struct {
    int fd;
    uint64_t size;
} BlockDeviceBackFile;

static int 
drive_read_sector (void *opaque, uint64_t sector_num, uint8_t *buf, int n) {
    BlockDeviceBackFile *bf;
    bf = opaque;
    uint64_t offset;
    offset = sector_num * SECTOR_SIZE;
    if (offset + n * SECTOR_SIZE > bf->size)
        return -1;
    if (pread (bf->fd, buf, n * SECTOR_SIZE, offset) < 0)
        return -1;
    return 0;
}

static int
drive_write_sector (void *opaque, uint64_t sector_num, const uint8_t *buf, int n) {
    BlockDeviceBackFile *bf;
    bf = opaque;
    uint64_t offset;
    offset = sector_num * SECTOR_SIZE;
    if (offset + n * SECTOR_SIZE > bf->size)
        return -1;
    if (pwrite (bf->fd, buf, n * SECTOR_SIZE, offset) < 0)
        return -1;
    return 0;
}

static int64_t
drive_get_sector_count (void *opaque) {
    BlockDeviceBackFile *bf;
    bf = opaque;
    return bf->size / SECTOR_SIZE;
}

BlockDevice *
drive_init (const char * f) {
    int fd, sz;
    BlockDevice *s;
    BlockDeviceBackFile *bf;
    if (access (f, F_OK | R_OK | W_OK) < 0)
        return NULL;
    s = mallocz (sizeof (*s));
    bf = mallocz (sizeof (*bf));
    fd = open (f, O_RDWR);
    sz = lseek (fd, 0, SEEK_END);
    bf->fd = fd;
    bf->size = sz;
    s->opaque = bf;
    s->get_sector_count = drive_get_sector_count;
    s->read_sector = drive_read_sector;
    s->write_sector = drive_write_sector;
    return s;
}

static void
virt_machine_run (RISCVMachine *m) {
    /* handle input */
    int stdin_fd;
    struct pollfd p;
    STDIODevice *s = m->cs->opaque;
    stdin_fd = s->stdin_fd;
    BOOL timer_intr;
    timer_intr = TRUE;
    if (virtio_console_can_write_data (m->console)) {
        p.fd = stdin_fd;
        p.events = POLLIN;
        if (poll (&p, 1, 0) == 1 && (p.events & POLLIN) != 0) {
            uint8_t buf[128];
            int ret, len;
            len = virtio_console_get_write_len (m->console);
            len = min_int (len, sizeof (buf));
            ret = m->cs->read_data(m->cs->opaque, buf, len);
            virtio_console_write_data (m->console, buf, ret);
            timer_intr = FALSE;
        }
    }
    
    /* handle timer interrupt */
    if (timer_intr)
        raise_timer_interrupt (m);
    riscv_cpu_interp (m->cpu_state, MAX_EXEC_CYCLE);
}

static BOOL
check_virt_machine_params (VirtMachineParams *p) {
    if (p->console != NULL)
    if (p->drive != NULL)
    if (access (p->kernel, F_OK | R_OK) == 0)
    if (p->ram_size > 0)
    return TRUE;
    return FALSE;
}

int
main (int argc, char ** argv) {
    if (argc != 2) {
        printf ("usage: %s kernel\n", argv[0]);
        exit (EXIT_SUCCESS);
    }
    VirtMachineParams p;
    term_init (TRUE);
    p.console = console_init ();
    p.drive = drive_init (DEFAULT_DRIVE);
    p.ram_size = MAX_RAM_SIZE;
    p.kernel = argv[1];
    RISCVMachine *m;
    if (check_virt_machine_params (&p) == FALSE) {
        printf ("invalid parameters.\n");
        exit (EXIT_FAILURE);
    }
    if ((m = machine_init (&p)) == NULL) {
        printf ("failed to init machine.\n");
        exit (EXIT_FAILURE);
    }
    for (;;)
        virt_machine_run (m);
    return EXIT_SUCCESS;
}